#include <common/HeaderSpace.h>
#include <folly/Conv.h>
#include <folly/gen/Base.h>
#include <folly/logging/xlog.h>
#include <pipeline/ActionAggregator.h>
#include <range/v3/view/group_by.hpp>
#include <numeric>
#include <execution>

namespace flashpp {

ActionAggregator::ActionAggregator(Pipeline& pipeline) {
  addFiberTask([&pipeline, this]() mutable noexcept {
    while (1) {
      std::shared_ptr<AtomicOverwrite> ao;
      auto ticket = pipeline.blockingReadStage<2>(ao);
      auto out = std::make_shared<AtomicOverwrite>();
      out->device = ao->device;
      process(*ao, *out);
      pipeline.blockingWriteStage<2>(ticket, std::move(out));
    }
  });
}

void
ActionAggregator::process(AtomicOverwrite& ao, AtomicOverwrite& out) {

  std::map<Port*, std::vector<HeaderSpace>> grouped;
  for (auto& o : ao.overwrites) {
    grouped[o.second].push_back(std::move(o.first));
    // grouped[o.second] = grouped[o.second] | o.first;
  }

  for (auto& [k, v] : grouped) {
    auto res = std::reduce(std::execution::par_unseq, v.begin(), v.end(), HeaderSpace::hsEmpty(), [](auto x, auto y){return x | y;});
    out.overwrites.push_back(std::make_pair(res, k));
    // out.overwrites.push_back(std::make_pair(HeaderSpace::orAll(v), k));
  }
}

} // namespace flashpp